﻿//#define VERBOSE

using System;
using System.Collections.Generic;
using System.Diagnostics;
using RayDen.Library.Components.Surface;
using RayDen.Library.Core;
using RayDen.Library.Core.Primitives;
using RayDen.Library.Entity.Scene;
using RayDen.RayEngine.Core.Interface;
using RayDen.RayEngine.Core.Types;
using RayDen.RayEngine.Engines.PathTracer;

namespace RayDen.RayEngine.Engines.MCMC
{
    // ReSharper disable RedundantThisQualifier
    
    public enum MCMCSamplerState
    {
        Initialized,
        SmallStep,
        LargeStep,
    }

    public class MLTSample
    {
        public Sample sample;
        public float Weight;
        public RgbSpectrum Contribution;
        public int ModifyTime;
    }



    public class MCMCPathTracerPath : PathSamplerBase
    {
        private string id = Guid.NewGuid().ToString();

        public string SamplerID
        {
            get { return string.Format("PathSampler_{0}", id); }
        }

        internal RgbSpectrum Throughput, TentativeContrib;

        public Sample TentativeSample;
        public MCMCSamplerState SamplerState;
        public PathTracerPathState PathState;
        public float xImportance;
        private RgbSpectrum xContribution;

        private Stack<MLTSample> sampleCache;

        protected int depth, tracedShadowRayCount, mutationsCount, rejectCount;
        protected internal int time;
        protected internal float pathWeight, TentativeWeight;
        protected bool specularBounce;

        public UInt64 PathExpression;

        protected RayEngineScene scene;


        public int MaxRaysPerPath
        {
            get;
            set;
        }

        struct ShadowRayInfo
        {
            public float pdf;
            public RgbSpectrum color;
            public RayData shadowRay;
            public int currentShadowRayIndex;
        }
        private ShadowRayInfo[] secRays;
        private MCMCSampler sampler;
        private bool mutate;


        public MCMCPathTracerPath()
        {
            sampleCache = new Stack<MLTSample>(10);
        }


        public void Push()
        {
            this.sampleCache.Push(new MLTSample()
                {
                    sample = new Sample(this.Sample),
                    Contribution = this.Radiance,
                    ModifyTime = this.time,
                    Weight = pathWeight
                });
        }


        public bool Pop()
        {
            if (this.sampleCache.Count > 0)
            {
                var sample = sampleCache.Pop();
                this.Sample = new Sample(sample.sample);
                this.pathWeight = sample.Weight;
                this.Radiance = sample.Contribution;
                return true;
            }
            return false;
        }

        public override void InitPath(IPathProcessor buffer)
        {
            base.InitPath(buffer);
            this.scene = pathIntegrator.Scene;
            this.Radiance = new RgbSpectrum(0f);
            this.Throughput = new RgbSpectrum(1f);
            this.PathState = PathTracerPathState.EyeVertex;
            this.secRays = new ShadowRayInfo[scene.ShadowRayCount];
            this.sampler = (MCMCSampler)pathIntegrator.Sampler;
            switch (this.SamplerState)
            {
                case MCMCSamplerState.LargeStep:
                case MCMCSamplerState.Initialized:
                    time = 0;
                    this.mutationsCount = 0;
                    this.rejectCount = 0;
                    //this.sampleCache.Clear();
                    this.Sample = pathIntegrator.Sampler.GetSample(null);
                    break;
                case MCMCSamplerState.SmallStep:
                    {
                        //if (rejectCount > 0)
                        //{
                        //    Print("Rejects : " + rejectCount);
                        //}
                        this.time++;
                        this.Sample = this.sampler.SmallStep(this.Sample); ;
                    }
                    break;
            }

            IRay ray;
            pathIntegrator.Scene.Camera.GetRay(Sample.imageX, Sample.imageY, out ray);
            this.PathRay = (RayData)ray;
            this.RayIndex = -1;
            this.pathWeight = 1.0f;

            this.tracedShadowRayCount = 0;
            this.depth = 0;
            this.specularBounce = true;
        }

        private void SwitchToSmallStep(float imp, float pathW, ref RgbSpectrum contribution)
        {
            this.xImportance = imp;
            this.xContribution = contribution;
            this.TentativeSample = new Sample(this.Sample);
            this.TentativeWeight = pathW;
            Push();
            this.SamplerState = MCMCSamplerState.SmallStep;
        }

        public override void Splat(SampleBuffer sampleBuffer)
        {
            //Check tentative prob
            //Add contribution
            //Change state
            // Change sample state in init by ChangedState

            var yImportance = sampler.EvalImportance(ref Throughput, pathWeight);
            float sampleWeight = yImportance;
            if (yImportance<=0f)
                goto @init;

            var yContribution = Radiance / yImportance;
            bool accept;
            float a = 0f;
            if (!sampler.ImportanceReady)
            {
                sampler.Add_B(ref yContribution);
                sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref Radiance);
                this.InitPath(this.pathIntegrator);
                return;
            }

            if (yImportance > 0f)
            {
                if ((SamplerState != MCMCSamplerState.SmallStep || xImportance <= 0f || mutate))
                {
                    SwitchToSmallStep(yImportance, pathWeight, ref yContribution);
                    goto init;
                    //Print(string.Format("{0} changed state to {1}", this.SamplerID, this.SamplerState));
                }
                else if (SamplerState == MCMCSamplerState.SmallStep)
                {
                    a = Math.Min(1f, (yImportance/xImportance));
                    //contrib = (contrib / sampler.Importance) / sampler.MaxMutations;
                    sampleWeight = a;
                    this.mutationsCount++;

                    if (sampler.NextFloat() < a)
                    {
                        xImportance = yImportance;
                        this.xContribution = yContribution;
                        //this.TentativeSample = new Sample(this.Sample);
                        Push();
                        accept = true;
                        //this.pathWeight = a;
                    }
                    else
                    {
                        if (!this.Pop())
                        {
                            //Print("Chain restart - stack empty - mutations " + mutationsCount);
                            //sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref contrib, 1f-a);
                            this.Sample = new Sample(this.TentativeSample);
                            rejectCount++;
                            this.pathWeight = this.TentativeWeight;
                        }

                        //Splat = false;
                        //return;
                    }
                    /*
                   
                         * */
                    //var newSample = sampler.GetNextSample(this.Sample, I, oldI, 1.0f, ref contrib, out outC, out newW, false);
                }

               

             
            }
            else if (xImportance >= 0f && this.SamplerState == MCMCSamplerState.SmallStep)
            {
                if (!this.Pop())
                {
                    if (this.TentativeSample != null)
                    {
                        //Print("Chain restart - stack empty - mutations " + mutationsCount);
                        //sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref contrib, 1f-a);
                        this.Sample = new Sample(this.TentativeSample);
                        rejectCount++;
                        this.pathWeight = this.TentativeWeight;
                    }
                    else
                    {
                        this.SamplerState = MCMCSamplerState.LargeStep;
                    }
                }

            }

            if (!xContribution.IsBlack())
            {
                //sampleWeight *= 1f/pathWeight;
                //contrib *= (1f/sampler.MaxMutations);
                //contrib *= new RgbSpectrum((float)this.sampler.B[0],(float)this.sampler.B[1],(float)this.sampler.B[2]);
                var contribution = xContribution;
                /*
                if (accept)
                {
                    contribution *= sampleWeight*sampler.Importance;
                }
                else
                {
                    contribution *= (1f - sampleWeight) * sampler.Importance;
                }
                */

                sampleBuffer.SplatSample(this.Sample.imageX, this.Sample.imageY, ref contribution);
            }

            if ((mutationsCount >= sampler.MaxMutations) & (sampler.NextFloat() > (1f - sampler.LargeStepProb)))//time > mutationsCount || 
            {
                this.SamplerState = MCMCSamplerState.LargeStep;
                this.time = 0;
                //Print(string.Format("{0} changed state to {1}", this.SamplerID, this.SamplerState));
                //Print("Chain restart - max mutations");
            }
            @init:
           
            this.InitPath(this.pathIntegrator);
        }

        public override bool FillRayBuffer(RayBuffer rayBuffer)
        {
            var leftSpace = rayBuffer.LeftSpace();
            if (((PathState == PathTracerPathState.EyeVertex) && (1 > leftSpace)) ||
            ((PathState == PathTracerPathState.ShadowRaysOnly) && (tracedShadowRayCount > leftSpace)) ||
            ((PathState == PathTracerPathState.NextVertex) && (tracedShadowRayCount + 1 > leftSpace)))
                return false;
            if (PathState != PathTracerPathState.ShadowRaysOnly)
                RayIndex = rayBuffer.AddRay(PathRay);
            if (PathState == PathTracerPathState.NextVertex || PathState == PathTracerPathState.ShadowRaysOnly)
            {
                for (int i = 0; i < tracedShadowRayCount; ++i)
                    secRays[i].currentShadowRayIndex = rayBuffer.AddRay(secRays[i].shadowRay);
            }
            return true;
        }


        public override void Advance(RayBuffer rayBuffer, SampleBuffer consumer)
        {
            int currentTriangleIndex = 0;
            SurfaceIntersectionData hitInfo = null;
#if VERBOSE

            try
            {
#endif

                base.Advance(rayBuffer, consumer);

                RayHit rayHit = rayBuffer.rayHits[RayIndex];

                if (((PathState == PathTracerPathState.ShadowRaysOnly) || (PathState == PathTracerPathState.NextVertex)) &&
                    (tracedShadowRayCount > 0))
                {
                    for (int i = 0; i < tracedShadowRayCount; ++i)
                    {
                        RayHit shadowRayHit = rayBuffer.rayHits[secRays[i].currentShadowRayIndex];
                        RgbSpectrum attenuation;
                        if (this.ShadowRayTest(ref shadowRayHit, out attenuation))
                        {
                            //                            Radiance.MADD()
                            mutate = true;
                            Radiance += attenuation * ((secRays[i].color));
                            pathWeight *= secRays[i].pdf;
                        }
                    }
                    Splat(consumer);

                    return;
                }

                depth++;
                bool missed = rayHit.Index == 0xffffffffu;
                if (missed || PathState == PathTracerPathState.ShadowRaysOnly || depth > scene.MaxPathDepth)
                {
                    if (specularBounce && missed)
                    {
                        Radiance += this.SampleEnvironment(PathRay.Dir) * Throughput;
                    }
                    Splat(consumer);

                    return;
                }

                // Something was hit
                hitInfo = SurfaceSampler.GetIntersection(ref PathRay, ref rayHit);
                currentTriangleIndex = (int)rayHit.Index;

                if (hitInfo == null)
                {
                    Debugger.Break();
                }
                //If Hit light)
                if (hitInfo.IsLight)
                {
                    if (specularBounce)
                    {
                        var lt = scene.GetLightByIndex(currentTriangleIndex);
                        if (lt != null)
                        {
                            var le = (RgbSpectrum)(lt.Le(ref PathRay.Dir));
                            Radiance += Throughput * le;
                            mutate = true;
                        }
                    }
                    Splat(consumer);

                    return;
                }

                var hitPoint = PathRay.Point(rayHit.Distance);

                Vector wo = -PathRay.Dir;
                tracedShadowRayCount = 0;
                if (hitInfo.MMaterial.IsDiffuse())
                {
                    float lightStrategyPdf = 1f;// scene.ShadowRayCount / (float)scene.Lights.Length;
                    RgbSpectrum lightTroughtput = Throughput * hitInfo.Color;
                    for (int i = 0; i < scene.ShadowRayCount; ++i)
                    {
                        int currentLightIndex = scene.SampleLights(Sample.GetLazyValue(depth * 0));
                        var light = scene.Lights[currentLightIndex];

                        var ls = new LightSample();
                        light.EvaluateShadow(ref hitPoint, ref hitInfo.ShadingNormal, Sample.GetLazyValue(depth * 1), Sample.GetLazyValue(depth * 2), Sample.GetLazyValue(depth * 3), ref ls);
                        if (ls.Pdf <= 0f)
                        {
                            continue;
                        }

                        secRays[tracedShadowRayCount].color = (RgbSpectrum)(ls.Spectrum);
                        secRays[tracedShadowRayCount].pdf = ls.Pdf;
                        secRays[tracedShadowRayCount].shadowRay = ls.LightRay;

                        Vector lwi = secRays[tracedShadowRayCount].shadowRay.Dir;
                        RgbSpectrum fs;
                        hitInfo.MMaterial.f(
                            ref secRays[tracedShadowRayCount].shadowRay.Dir,
                            ref wo, ref hitInfo.ShadingNormal, ref Throughput, out fs, types: BrdfType.Diffuse);
                        secRays[tracedShadowRayCount].color *= lightTroughtput *
                                                               Vector.AbsDot(ref hitInfo.ShadingNormal, ref lwi) *
                                                               fs;
                        if (!secRays[tracedShadowRayCount].color.IsBlack())
                        {
                            secRays[tracedShadowRayCount].pdf *= lightStrategyPdf;
                            tracedShadowRayCount++;
                        }
                    }
                }

                float fPdf;
                Vector wi;

                RgbSpectrum f = hitInfo.MMaterial.Sample_f(ref wo, out wi, ref hitInfo.Normal, ref hitInfo.ShadingNormal, ref Throughput,
                                                          Sample.GetLazyValue(depth * 4), Sample.GetLazyValue(depth * 5),
                                                          Sample.GetLazyValue(depth * 6), ref hitInfo.TextureData, 
                                                          out fPdf, out specularBounce);

                if ((fPdf <= 0.0f) || f.IsBlack())
                {
                    if (tracedShadowRayCount > 0)
                        PathState = PathTracerPathState.ShadowRaysOnly;
                    else
                    {
                        Splat(consumer);

                    }
                    return;
                }
                pathWeight *= fPdf;
                Throughput *= f / fPdf;
                if (!mutate)
                    mutate = Throughput.Filter() > Sample.GetLazyValue(depth*7);

                if (depth > scene.MaxPathDepth)
                {
                    float prob = Math.Max(Throughput.Filter(), scene.RussianRuletteImportanceCap);
                    if (prob >= Sample.GetLazyValue(depth * 8))
                    {
                        Throughput /= prob;
                        pathWeight *= prob;
                    }
                    else
                    {
                        if (tracedShadowRayCount > 0)
                            PathState = PathTracerPathState.ShadowRaysOnly;
                        else
                        {
                            Splat(consumer);
                        }

                        return;
                    }
                }

                PathRay.Org = hitPoint;
                PathRay.Dir = wi.Normalize();
                PathState = PathTracerPathState.NextVertex;

#if VERBOSE
            }
            catch (Exception ex)
            {
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Advance path exception");
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine("Error triangle {0}", currentTriangleIndex);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.Message);
                RayDen.Library.Components.SystemComponents.Tracer.TraceLine(ex.StackTrace);

            }
#endif
        }


        private bool ShadowRayTest(ref RayHit shadowRayHit, out RgbSpectrum attenuation)
        {
            var hit = shadowRayHit.Index == 0xffffffffu || scene.IsLight((int)shadowRayHit.Index);
            attenuation = new RgbSpectrum(1f);
            if (hit) return true;
            var mesh = scene.GetMeshByTriangleIndex((int)shadowRayHit.Index);
            var mat =
                //SurfaceMaterials.CreateMaterial(scene.MaterialProvider.Get(mesh.MaterialName));
                scene.MatLib.GetSurfMat(mesh != null ? mesh.MaterialName : "default");

            if (mat.AlphaTexture != null)
            {
                attenuation = (RgbSpectrum) (mat.AlphaTexture as ConstTextureInfo).Color;
                hit = true;
                //mat.MaterialData.Kt.Filter() < Sample.GetLazyValue();
            }

            return hit;
        }

        private RgbSpectrum SampleEnvironment(Vector vector)
        {
            return this.scene.SampleEnvironment(-vector);
        }


    }


    // ReSharper restore RedundantThisQualifier
}